"""
supported models:
- resnet
- bigtransfer
- masked autoencoder
- contrastive learning (moco)
- generative learning (iBot)
"""
from .resnet import ResnetFeatureExtractor
from .big_transfer import BiTFeatureExtractor
from .masked_autoencoders import MAEFeatureExtractor
from .CL_Extractor import build_CL_extractor
from .GL_Extractor import build_GL_extractor


SUPPORTED_NAMES = [
    'resnet18',
    'resnet34',
    'resnet50',
    'resnet101',
    'BiT-M-R50x1',
    'BiT-M-R50x3',
    'MAE_vit_base_patch16',
    'CL_vit_small_patch16',
    'GL_vit_small_patch16',
]


def get_feature_extractor(config):
    model_name = config['NAME']
    assert model_name in SUPPORTED_NAMES, "unsupported model type"

    if model_name.startswith('resnet'):
        model = ResnetFeatureExtractor(model_name)
    elif model_name.startswith('BiT'):
        model_name = model_name.replace('BiT_', '')
        model = BiTFeatureExtractor(model_name)
    elif model_name.startswith('MAE'):
        model_name = model_name.replace('MAE_', '')
        model = MAEFeatureExtractor(model_name)
    elif model_name.startswith('CL'):
        model = build_CL_extractor()
    elif model_name.startswith('GL'):
        model = build_GL_extractor()
    else:
        raise ValueError('unsupported model type')
    return model
